from langdev_helper.prompt.writer import *
from langdev_helper.util.code import *
from langdev_core.llms.qianfan import call_llm
from langchain_core.tools import tool
from langchain.tools.render import render_text_description
from langchain_core.prompts import PromptTemplate
import json




def do_llm():
    prompt_user = prompt_outline + prompt_outline_format
    print(prompt_user)

    return call_llm(prompt_user, system_prompt)

def parse_outline(res):
    code = extract_code(res, 'json')
    # print(code)
    data = json.loads(code)
    # print(data)
    return data['小说设定'], data['小说大纲']

def write_outline():
    return parse_outline(do_llm())

def parse_chapter(chapter):
    return chapter['标题'], chapter['内容']

def write_chapter(chapter):
    title, content = parse_chapter(chapter)
    prompt_outline_chapter_template = PromptTemplate.from_template(prompt_outline_chapter)
    outline_chapter_pt = prompt_outline_chapter_template.format(title=title, content=content)
    print("========outline_chapter_pt")
    print(outline_chapter_pt)
    print("========res")
    res = call_llm(outline_chapter_pt, system_prompt)
    print(res)

    prompt_outline_detail_tpl =  PromptTemplate.from_template(prompt_outline_detail)
    prompt_outline_detail_pt = prompt_outline_detail_tpl.format(chapter_outline=res)
    print("========prompt_outline_detail_pt")
    print(prompt_outline_detail_pt)
    print("========res")
    res = call_llm(prompt_outline_detail_pt, system_prompt)
    print(res)


    prompt_chapter_content_tpl =  PromptTemplate.from_template(prompt_chapter_content)
    prompt_chapter_content_pt = prompt_chapter_content_tpl.format(detail_outline=res, word_count=1000)
    print("========prompt_chapter_content_pt")
    print(prompt_chapter_content_pt)
    print("========res")
    res = call_llm(prompt_chapter_content_pt, system_prompt)
    print(res)


def write_chapter_outline():
    preset , outline= write_outline()

    for chapter in outline:
        name, content = parse_chapter(chapter)
        print(name, '==>', content)
    write_chapter(outline[0])

write_chapter_outline()
